import community
from shell.knowledge_graph_mining.prepare_graph_songs import GraphSongs
import json


def louvain_heuristices(graph):
    partition = community.best_partition(graph)
    return partition


def modularity(graph, partition):
    return community.modularity(partition, graph)


def cluster():
    json_str = json.load(open('../../data/data.json'))
    gs_obj = GraphSongs(json_str)
    graph = gs_obj.prepare_graph(list(json_str.keys()))
    partition = louvain_heuristices(graph)
    print('partition:', partition)
    print('modularity:', modularity(graph, partition))


if __name__ == '__main__':
    cluster()
